"""
增强的记忆管理类
基于现有MemoryClass进行优化，使用传统记忆系统
注意：异步存储功能已移除，避免事件循环问题
"""

import json
from typing import Dict, Any, List, Optional
from datetime import datetime, timedelta
import logging

from langchain.memory import ConversationBufferMemory
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# QueryRewriter导入已移除 - 统一在LangGraph Supervisor中使用

logger = logging.getLogger(__name__)

class EnhancedMemoryClass:
    """增强的记忆管理类（简化版，移除异步存储）"""

    def __init__(
        self,
        memoryKey: str = "chat_history",
        model: str = None,
        config: Dict[str, Any] = None
    ):
        """初始化增强记忆管理（简化版）"""
        import os

        self.memoryKey = memoryKey
        self.model_name = model or os.getenv("BASE_MODEL", "gpt-4")
        self.config = config or {}
        
        # 初始化模型
        self.chatmodel = ChatOpenAI(
            model=self.model_name, 
            base_url=os.getenv("OPENAI_API_BASE"),
            api_key=os.getenv("OPENAI_API_KEY")
        )
        
        # 传统记忆系统
        self.memory_dict = {}

        # Redis配置
        self.redis_url = os.getenv("REDIS_URL", "redis://localhost:6380/0")

        # 查询重写器已移除 - 统一在LangGraph Supervisor中处理
        # 避免重复实现，降低系统复杂度
        self.query_rewriter = None

        logger.info("Enhanced Memory Class initialized (simplified mode)")
    

    
    def store_conversation_turn(
        self,
        user_message: str,
        ai_response: str,
        user_id: str,
        session_id: str,
        metadata: Dict[str, Any] = None
    ) -> bool:
        """存储对话轮次（使用传统存储）"""
        try:
            # 使用传统存储
            return self._store_conversation_traditional(
                user_message, ai_response, session_id
            )

        except Exception as e:
            logger.error(f"Failed to store conversation turn: {e}")
            return False
    
    def _store_conversation_traditional(
        self, 
        user_message: str, 
        ai_response: str, 
        session_id: str
    ) -> bool:
        """传统方式存储对话（降级方案）"""
        try:
            chat_history = RedisChatMessageHistory(
                session_id=session_id, 
                url=self.redis_url
            )
            
            from langchain_core.messages import HumanMessage, AIMessage
            
            chat_history.add_message(HumanMessage(content=user_message))
            chat_history.add_message(AIMessage(content=ai_response))
            
            return True
            
        except Exception as e:
            logger.error(f"Failed to store conversation traditionally: {e}")
            return False
    
    def get_enhanced_context(
        self,
        current_query: str,
        user_id: str,
        session_id: str,
        include_rewrite: bool = True
    ) -> Dict[str, Any]:
        """获取增强的上下文信息（简化版 - 查询重写已移至LangGraph Supervisor）"""
        try:
            # 简化的上下文，不再包含查询重写逻辑
            # 查询重写统一在LangGraph Supervisor中处理，避免重复计算
            context = {
                "original_query": current_query,
                "rewritten_query": current_query,  # 保持接口兼容性
                "conversation_history": self._get_traditional_history(session_id),
                "relevant_memories": [],
                "user_profile": {},
                "rewrite_info": None  # 重写信息由Supervisor提供
            }

            logger.debug(f"Enhanced context prepared for query: {current_query[:50]}...")
            return context

        except Exception as e:
            logger.error(f"Failed to get enhanced context: {e}")
            # 返回基础上下文
            return {
                "original_query": current_query,
                "rewritten_query": current_query,
                "conversation_history": self._get_traditional_history(session_id),
                "relevant_memories": [],
                "user_profile": {},
                "rewrite_info": None
            }
    
    def _get_traditional_history(self, session_id: str) -> List[Dict[str, str]]:
        """获取传统方式的对话历史"""
        try:
            chat_history = RedisChatMessageHistory(
                session_id=session_id, 
                url=self.redis_url
            )
            
            history = []
            for message in chat_history.messages[-20:]:  # 最近20条消息
                role = "user" if message.type == "human" else "assistant"
                history.append({
                    "role": role,
                    "content": message.content,
                    "timestamp": datetime.now().isoformat()
                })
            
            return history
            
        except Exception as e:
            logger.error(f"Failed to get traditional history: {e}")
            return []
    

    
    # 向后兼容的方法
    def get_memory(self, session_id: str = "session1"):
        """获取记忆（向后兼容）"""
        try:
            chat_message_history = RedisChatMessageHistory(
                session_id=session_id, url=self.redis_url
            )
            
            # 检查是否需要摘要（使用传统方式）
            store_message = chat_message_history.messages
            if hasattr(self.config, 'LONG_TERM_SUMMARY_THRESHOLD') and len(store_message) > self.config.LONG_TERM_SUMMARY_THRESHOLD:
                # 使用传统摘要方式
                str_message = ""
                for message in store_message:
                    str_message += f"{type(message).__name__}: {message.content}\n"
                summary = self.summary_chain(str_message)
                chat_message_history.clear()
                chat_message_history.add_message(summary)
            
            return chat_message_history
            
        except Exception as e:
            logger.error(f"Error in get_memory: {e}")
            return None
    
    def set_memory(self, session_id: str = "session1"):
        """设置记忆（向后兼容）"""
        chat_memory = self.get_memory(session_id)
        if chat_memory is None:
            chat_memory = RedisChatMessageHistory(
                session_id=session_id, url=self.redis_url
            )
        
        self.memory = ConversationBufferMemory(
            llm=self.chatmodel,
            human_prefix="user",
            ai_prefix="小智",
            memory_key=self.memoryKey,
            output_key="output",
            return_messages=True,
            max_token_limit=3000,
            chat_memory=chat_memory
        )
        
        return self.memory
    
    def summary_chain(self, store_message: str):
        """摘要链（向后兼容）"""
        try:
            from dingding.Prompt import PromptClass
            
            prompt_instance = PromptClass()
            SystemPrompt = prompt_instance.system_prompt.format(feelScore=5, who_you_are="")
            
            prompt = ChatPromptTemplate.from_messages([
                ("system", SystemPrompt + "\n这是一段你和用户的对话记忆，对其进行总结摘要，摘要使用第一人称'我'，并且提取其中的关键信息，以如下格式返回：\n 总结摘要 | 过去对话关键信息\n例如 用户张三问候我好，我礼貌回复，然后他问我langchain的向量库信息，我回答了他今年的问题，然后他又问了比特币价格。|Langchain, 向量库,比特币价格"),
                ("user", "{input}")
            ])
            
            chain = prompt | self.chatmodel
            summary = chain.invoke({"input": store_message})
            return summary
            
        except Exception as e:
            logger.error(f"摘要出错: {e}")
            return None
